import random
import math

import torch
import torch.utils.data as Data

class RandBatchSampler(Data.Sampler):
    def __init__(self, labels_list, batch_size, subset_size=4, weights=None, cum_weights=None, replacement=False):

        # 计算类标签数与类标签有哪些
        self.cls_list = [] # 记录类别标签
        self.labels_list = [] # 记录与上列表中labels对应的数据索引有哪些
        tmp = {} # 暂时记录labels_list中的labels与其在cls_list中对应索引的映射关系
        for n, cls_labels in enumerate(labels_list):
            if cls_labels not in self.cls_list:
                self.cls_list.append(cls_labels)
                self.labels_list.append([])
                tmp[str(cls_labels)] = len(self.cls_list) - 1
            self.labels_list[tmp[str(cls_labels)]].append(n)
        self.classes = len(self.cls_list)

        func = lambda elem : self.cls_list[elem]
        tmp_ids_list = list(range(self.classes))
        # tmp_ids_list 的每个元素的值和位置代表着当前索引为该值的labels_list
        # cls_list元素在按照类别符号从小到大排序后应处于的位置
        tmp_ids_list.sort(key=func) 
        self.cls_list.sort()
        for n in range(self.classes):
            self.labels_list[n].append(n)
        func = lambda elem : tmp_ids_list.index(elem[-1])
        self.labels_list.sort(key=func)

        for n in range(self.classes):
            self.labels_list[n].pop()

        # 计算每类的样本数
        self.cls_samples_num = []
        for cls_labels in self.labels_list:
            self.cls_samples_num.append(len(cls_labels))

        self.sample_num = len(labels_list)
        self.replacement = replacement # 采样是否重复

        # batch中子集大小，为0时关闭该功能
        self.subset_size = subset_size
        self.batch_size = batch_size

        if self.subset_size != 0 and self.batch_size % self.subset_size !=0:
            raise ValueError("batch_size % subset_size !=0")

        # 如果 self.batch_size == 0 即 没有开batch，此时self.subset无效，迭代器的返回值并非列表，此时self.cls_num=1
        # 如果 self.batch_size != 0 即 batch模式，迭代器的返回值为列表，此时self.cls_num>=1
        # 在上一种情况下， 如果self.subset == 0 or 1 即 没开subset模式，batch中的sample是随机选择（一个一个地选）
        #               如果self.subset > 1 即打开subset模式， batch中的sample是一次性选出子集数量个数（数个数个地选）
        if self.batch_size != 0 and self.subset_size == 0:
            self.subset_size = 1
        self.cls_num = self.batch_size // self.subset_size if self.batch_size != 0 else 0

        # 转成权重的累加值
        if cum_weights != None:
            self.weights = cum_weights.copy()
        elif weights != None:
            self.weights = weights.copy()
            for i in range(len(self.weights)):
                tmp_weights = self.weights[i]
                if i != 0:
                    tmp_weights += self.weights[i-1]
                    self.weights[i] = tmp_weights
        # 当weights为None时默认每个样本采样概率都是一样的
        # 所以self.weights应当对应每个类别数量占比的累加
        else:
            self.weights = []
            for n in range(self.classes):
                if n == 0:
                    self.weights.append(self.cls_samples_num[n])
                else:
                    self.weights.append(self.cls_samples_num[n] + self.weights[n-1])

            for n in range(self.classes):
                self.weights[n] = self.weights[n] / self.weights[-1]

        if self.weights != None and len(self.weights) != self.classes:
            raise ValueError("wrong length of weights or cum_weights")
        if self.weights[-1] != 1.0:
            raise ValueError("last weights is not 1.0 but %.2f" % self.weights[-1])
        # print(self.weights)

    def __iter__(self):
        iter_num = self.__len__()

        # 有重复采样
        if self.replacement:
            for iter_i in range(iter_num):
                random.seed()

                # 不论是否打开batch模式，统一先采样类别
                cls_num = self.cls_num
                random_cls = random.choices(list(range(self.classes)), cum_weights=self.weights, k=cls_num)
                # 对每类分别采样样本
                sample_num = 1 if self.batch_size == 0 else self.subset_size # 每类中采样几个样本
                samples_ids = []
                for cls in random_cls:
                    samples_count = self.cls_samples_num[cls]

                    # 如果类别中样本个数少于采样个数，则会对一些样本重复采样
                    if sample_num > samples_count:
                        tmp1, tmp2 = sample_num // samples_count, sample_num % samples_count
                        sample_list = [] # 进行元素个数次采样， 每次采样对应元素个样本
                        for i in range(tmp1):
                            sample_list.append(samples_count)
                        sample_list.append(tmp2)

                        samples_ids_incls = []
                        for nums in sample_list:
                            samples_ids_incls += random.choices(list(range(samples_count)), k=nums)

                    else:
                        samples_ids_incls = random.choices(list(range(samples_count)), k=sample_num)

                    cls_sample_ids = []
                    for ids_ in samples_ids_incls:
                        cls_sample_ids.append(self.labels_list[cls][ids_])

                    samples_ids += cls_sample_ids

                if self.batch_size == 0:
                    yield samples_ids[0]
                else:
                    yield samples_ids
        # 无重复采样
        else:
            if self.batch_size == 0:
                cls_choices_sum = iter_num
            else:
                cls_batch_size = self.batch_size // self.subset_size
                cls_choices_sum = iter_num * cls_batch_size

            cls_choices_list = [] # 用于决定各类别的采样顺序
            for cls_ids in range(self.classes):
                if cls_ids == 0:
                    cls_choices_list += [cls_ids] \
                        * round(self.weights[cls_ids] * cls_choices_sum)
                else:
                    cls_choices_list += [cls_ids] \
                        * (round(self.weights[cls_ids] * cls_choices_sum) - \
                            round(self.weights[cls_ids-1] * cls_choices_sum))
            random.shuffle(cls_choices_list)

            cls_sample_choices_start = [0] * self.classes # 用于标记每类的样本在当前迭代时应当从某个位置开始采样
            for i in range(len(self.labels_list)):
                random.shuffle(self.labels_list[i])
                
            for i in range(iter_num):
                cur_cls_choices = cls_choices_list[cls_batch_size*i:cls_batch_size*(i+1)]
                cur_batch = []
                for j in cur_cls_choices:
                    start = cls_sample_choices_start[j]
                    end = start + self.subset_size
                    if end > len(self.labels_list[j]):
                        end_ptr = [end]
                        cur_batch += self.list_elem(self.labels_list[j], slice(start, end), end_ptr)
                        end = end_ptr[0]
                    else:
                        cur_batch += self.labels_list[j][start:end]
                    cls_sample_choices_start[j] = end

                yield cur_batch

    def __len__(self):
        if self.batch_size == 0:
            return self.sample_num
        else:
            return math.ceil(self.sample_num / self.batch_size)

    def get_classes(self):
        return self.classes

    def get_nums(self):
        return self.sample_num

    # 返回list_的ids_切片的元素，如果超出list_大小则用list_靠前的元素补充
    def list_elem(self, list_, ids_, end_=None):
        if type(ids_) != slice:
            raise TypeError("ids_ need type slice")
        if ids_.stop > len(list_):
            results = []
            end = ids_.stop
            start = ids_.start
            length = len(list_)
            while end > length:
                results += list_[start:]
                start = 0
                end -= length
            results += list_[start:end]
        else:
            results = list_[start:end]

        if end_ != None:
            end_[0] = end

        return results

if __name__ == '__main__':
    classes = 4
    labels_list = [0] * 50 + [1] * 20 + [2] * 40 + [3] * 30
    random.shuffle(labels_list)
    weights = []
    sum = 0
    for i in range(classes-1):
        weights.append(random.randint(10, 100 // classes - 1))
        sum += weights[i]
    weights.append(100 - sum)
    for i in range(classes):
        weights[i] /= 100
    # weights = [0.36, 0.14, 0.29, 0.21]
    RSampler = RandBatchSampler(labels_list, 48, weights=weights)

    # # 测试 random.choices()
    # cls_sum = [0] * classes
    # for i in range(100000):
    #     random.seed()
    #     # print(random.choices(list(range(4)), weights=weights, k=3))
    #     random_samples = random.choices(list(range(4)), weights=weights, k=3)
    #     for j in range(classes):
    #         cls_sum[j] += random_samples.count(j)
        
    # sum = 0
    # for j in range(classes):
    #     sum += cls_sum[j]
    # for j in range(classes):
    #     cls_sum[j] /= sum
    # print(cls_sum)


    cls_sum = [0] * classes
    samples_sum = [0] * len(labels_list)
    for j in range(1000):
        for i in RSampler:
            for ids in i:
                samples_sum[ids] += 1
                cls_sum[labels_list[ids]] += 1

    sum = 0
    for j in range(classes):
        sum += cls_sum[j]
    for j in range(classes):
        cls_sum[j] = round(cls_sum[j] / sum, 6)

    sum = 0
    for j in range(len(labels_list)):
        sum += samples_sum[j]
    for j in range(len(labels_list)):
        samples_sum[j] = round(samples_sum[j] / sum, 6)
    
    print(weights)
    print(cls_sum)
    print(samples_sum)
    